"""
author：石沙
date：2020-09-28
content：本模块来提取图片特征
"""

# 如下导入时为保证训练时的任务流能正常执行
import sys
from settings import MAIN_PATH, SRC_PATH
sys.path.extend([MAIN_PATH, SRC_PATH])


from PIL import Image
from PIL import ImageFile
import numpy as np
import torch
import pandas as pd
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from torch import unsqueeze
from torchvision.models import resnet34
import torch.nn as nn
import os
from settings import COVER_PATH
from site_packages.utils.models import ModelOp


ImageFile.LOAD_TRUNCATED_IMAGES = True
DEVICE = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")


class ImagePathFile:

    """用于生成记录图书封面图片id、路径信息的csv文件"""

    def __init__(self, root_dir, folders):
        self.root_dir = root_dir
        self.folders = folders

    def __call__(self):
        dfs = []
        csv_file = os.path.join(self.root_dir, 'cover_info.csv')
        for folder in self.folders:
            df_by_folder = pd.DataFrame(columns=['image_id', 'filename', 'path'])
            root, dirs, filename_by_folder = next(os.walk(os.path.join(self.root_dir, folder), topdown=False))
            df_by_folder['image_id'] = list(map(lambda filename: filename.split('.')[0], filename_by_folder))
            df_by_folder['filename'] = filename_by_folder
            df_by_folder['path'] = os.path.join(self.root_dir, folder)
            dfs.append(df_by_folder)
        df_all = pd.concat(dfs)
        df_all.to_csv(csv_file, index=None, encoding='utf-8')
        print('已生成用于记录图片路径信息的csv文件：{}'.format(csv_file))


class ChannelPadding:

    """对图像数据维度超出和维度不足情况进行减少和增加维度的操作"""

    def __init__(self, channel_num=3, idx_pos=0):
        self.channel_num = channel_num
        self.idx_pos = idx_pos

    def __call__(self, sample):
        # 维度缺少，应对黑白图像
        if len(sample.shape) == 2:
            return sample.unsqueeze(self.idx_pos).repeat(self.channel_num, 1, 1)

        # channel数超出，应对不足
        if len(sample.shape) == 3:
            this_channel_num = sample.shape[self.idx_pos]
            if this_channel_num == 1:
                return sample.repeat(self.channel_num, 1, 1)
            if this_channel_num > self.channel_num:
                return sample.narrow(self.idx_pos, 0, self.channel_num)

        return sample


class CoverDataset(Dataset):

    """生成书籍封面数据集"""

    def __init__(self, root_dir, csv_file='cover_info.csv', transform=None):
        """
        csv_file（string）：带注释的csv文件的路径。
        root_dir（string）：包含所有图像的目录。
        transform（callable， optional）：一个样本上的可用的可选变换
        """
        self.info = pd.read_csv(os.path.join(root_dir, csv_file))
        self.root_dir = root_dir
        self.transform = transform

    def __len__(self):
        return self.info.shape[0]

    def __getitem__(self, idx):
        image_path = os.path.join(*self.info.loc[idx, ['path', 'filename']])
        image = Image.open(image_path)

        if self.transform:
            sample = self.transform(image)
        return sample


class Net(nn.Module):
    """
    resnet34预训练模型
    """
    def __init__(self):
        super(Net, self).__init__()
        self.model = resnet34(pretrained=True, progress=False)

        # 奇怪的结构
        def save_output(module, inputs, outputs):
            self.buffer = outputs

        self.model.avgpool.register_forward_hook(save_output)

    def forward(self, x):
        self.model(x)
        return self.buffer


class ImageFeature:
    """
    利用预训练模型抽取图像特征
    """

    def __init__(self, data_loader, batch_size, csv_file_path=None):
        self.net = Net().to(DEVICE)
        self.data_loader = data_loader
        self.batch_size = batch_size
        self.info = pd.read_csv(os.path.join(csv_file_path))

    def align(self, features):
        last = features[-1]
        features = features[:-1]
        np_features = np.array(features).reshape(-1, 512)
        return np.concatenate([np_features, last], axis=0)

    def extract(self):
        self.net.eval()
        features = []
        for i_batch, sample_batched in enumerate(self.data_loader):
            if i_batch % 100 == 0:
                print('当前查找位置：', i_batch)
            feature = self.net(sample_batched.to(DEVICE))
            feature = feature.squeeze()
            features.append(feature.cpu().detach().numpy())
        df_features = pd.DataFrame(self.align(features), index=self.info['image_id'].values)
        return df_features


if __name__ == '__main__':
    # 基本设置
    batch_size = 40

    # 图像路径记录
    folders = ['cover_0', 'cover_1', 'cover_2', 'cover_3', 'cover_4']
    path_recorder = ImagePathFile(COVER_PATH, folders)
    path_recorder()

    # 数据变换
    data_transformer = transforms.Compose([
            transforms.Resize([224, 224]), # 会将PIL的格式转化为pytorch的格式
            transforms.ToTensor(),
            ChannelPadding(channel_num=3, idx_pos=0),
            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                 std=[0.229, 0.224, 0.225]),

    ])

    # 数据集建立
    cover_data = CoverDataset(COVER_PATH,transform=data_transformer)
    data_loader = DataLoader(cover_data, batch_size=batch_size)

    # 特征提取与保存
    df_features = ImageFeature(data_loader, batch_size, csv_file_path=os.path.join(COVER_PATH, 'cover_info.csv')).extract()
    ModelOp.save(df_features, 'image_feature', is_model=False)
